"""
该代码主要实现了一个自定义的日志格式化类 ColoredFormatter，它使用 colorama 库为日志消息添加颜色，使日志在控制台输出时更加清晰和易于区分。
它可以为不同级别的日志（如 DEBUG, INFO, WARNING, ERROR, CRITICAL）的日志等级名称添加不同的颜色，同时可以为日志消息中的特定关键字添加不同的颜色。
"""

import logging
import pandas as pd
from colorama import Style, Fore, init
from logging.handlers import RotatingFileHandler
from rich.console import Console
from rich.text import Text
from rich.panel import Panel
from rich.table import Table
from rich.box import ROUNDED
from rich.align import Align
from art import text2art
import logging
import math
import re
import sys


#-------------------------------------------------------------------------------------
# 自定义日志格式化器
#-------------------------------------------------------------------------------------
console = Console()

class StreamToLogger:
    def __init__(self, logger, log_level=logging.INFO):
        self.logger = logger
        self.log_level = log_level
        self.linebuf = ''

    def write(self, buf):
        for line in buf.rstrip().splitlines():
            self.logger.log(self.log_level, line.rstrip())

    def flush(self):
        pass

class ColoredFormatter(logging.Formatter):
    LEVEL_COLORS = {
        logging.DEBUG: 'BLUE',
        logging.INFO: 'GREEN',
        logging.WARNING: 'YELLOW',
        logging.ERROR: 'RED',
        logging.CRITICAL: 'MAGENTA'
    }

    def __init__(self, fmt=None, datefmt=None, style='%', keywords=None, keyword_colors=None):
        super().__init__(fmt, datefmt, style)
        self.keywords = keywords if keywords else []
        self.keyword_colors = keyword_colors if keyword_colors else {}
        init()

    def format(self, record):
        # 根据日志等级设置不同的颜色
        try:
            # 获取当前日志等级对应的颜色
            color = getattr(__import__('colorama').Fore, self.LEVEL_COLORS[record.levelno])
            # 将日志等级名称添加颜色并添加颜色重置代码，避免颜色污染后续输出
            record.levelname = color + record.levelname + Style.RESET_ALL
        except AttributeError:
            # 若颜色不存在，不添加颜色
            pass
        # 突出显示消息中的关键字
        if self.keywords:
            record.msg = self.highlight_keywords(str(record.msg), self.keywords, self.keyword_colors)
        return super().format(record)

    def highlight_keywords(self, message, keywords, keyword_colors):
        for keyword in keywords:
            color = keyword_colors.get(keyword, Fore.YELLOW)
            highlighted_keyword = color + keyword + Style.RESET_ALL
            message = message.replace(keyword, highlighted_keyword)
        return message


    def reset_color(self):
        try:
            import colorama
            return colorama.Style.RESET_ALL
        except ImportError:
            return ''


class FileFormatterNoColor(logging.Formatter):
    def __init__(self, fmt=None, datefmt=None, style='%', keywords=None, keyword_colors=None):
        super().__init__(fmt, datefmt, style)
        self.keywords = keywords if keywords else []
        self.keyword_colors = keyword_colors if keyword_colors else {}

    def format(self, record):
        # 确保不包含任何颜色代码
        record.levelname = re.sub(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])', '', record.levelname)
        if self.keywords:
            record.msg = self.highlight_keywords(str(record.msg), self.keywords, self.keyword_colors)
        record.msg = re.sub(r'\s+', ' ', str(record.msg)).strip()
        record.msg = re.sub(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])', '', str(record.msg))
        return super().format(record)

    def highlight_keywords(self, message, keywords, keyword_colors):
        for keyword in keywords:
            color = keyword_colors.get(keyword, '')
            highlighted_keyword = keyword + color
            message = message.replace(keyword, highlighted_keyword)
        return message

def configure_logging(keywords=None, keyword_colors=None, handlers=None):
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    logger.propagate = False
    if not logger.handlers:
        # 控制台日志处理
        console_handler = logging.StreamHandler()
        formatter = ColoredFormatter(
            fmt='%(asctime)s < %(levelname)7s >: %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S',
            keywords=keywords,
            keyword_colors=keyword_colors
        )
        console_handler.setFormatter(formatter)
        logger.addHandler(console_handler)
        if handlers:
            for handler in handlers:
                if isinstance(handler, logging.FileHandler):
                    file_formatter = FileFormatterNoColor(
                        fmt='%(asctime)s < %(levelname)7s >: %(message)s',
                        datefmt='%Y-%m-%d %H:%M:%S'
                    )
                    handler.setFormatter(file_formatter)
                logger.addHandler(handler)
    return logger


# 调用日志部分
custom_keyword_colors = {
    'error': 'RED',
    'warning': 'YELLOW',
    'critical': 'MAGENTA',
    'successfully': 'GREEN'
}

file_handler = RotatingFileHandler('output.log', mode='a', maxBytes=1024 * 1024, backupCount=5)
logger = configure_logging(keywords=['error', 'warning', 'critical','successfully'],
                        keyword_colors=custom_keyword_colors,
                        handlers=[file_handler])


#-------------------------------------------------------------------------------------
# 自定义数据格式化打印
#-------------------------------------------------------------------------------------

def data_table(data, table_name=''):
    table = Table(header_style=table_name, show_header=True, box=ROUNDED)

    if isinstance(data, pd.DataFrame):
        columns = data.columns
        data_iter = data.iterrows()

    elif isinstance(data, list) and len(data) > 0 and isinstance(data[0], dict):
        columns = data[0].keys()
        data_iter = enumerate(data)

    elif isinstance(data, dict):
        columns = data.keys()
        data_iter = data.items()

    else:
        raise ValueError("Unsupported data format. Provide a pandas DataFrame or a list of dictionaries.")
    for key in columns:
        table.add_column(key, header_style="bright_cyan", justify="center")
    for _, item in data_iter:
        row = []
        for value in item.values() if isinstance(item, dict) else item:
            # Check if value is a float and is NaN, or if value is a string and equals 'nan'
            if (isinstance(value, float) and math.isnan(value)) or (isinstance(value, str) and value.lower() == 'nan'):
                row.append("N/A")
            else:
                row.append(str(value))
        table.add_row(*row)
    
    return table

def console_table(data, explain=''):
    """
    打印数据到控制台
    :param data: 要打印的数据，可以是一个列表、字典或 DataFrame
    :param explain: 对数据的解释或标题信息, 默认为None
    """
    try:
        if isinstance(data, pd.DataFrame):
                data_records = data
        elif isinstance(data, list):
            if len(data) > 0 and isinstance(data[0], dict):
                data_records = data
            else:
                # 将列表元素转换为字典
                new_data = []
                for item in data:
                    parts = item.split(',')
                    item_dict = {}
                    for part in parts:
                        key, value = part.split(':')
                        item_dict[key.strip()] = value.strip()
                    new_data.append(item_dict)
                data_records = new_data
        elif isinstance(data, dict):
            data_records = [data]
        elif isinstance(data, str):
            if data.endswith(('.xlsx', '.xls')):
                data_records = pd.read_excel(data)
            else:
                logger.error(f"{data} 文件类型错误，无法打印表格。")
                return None
        else:
            logger.error(f"不支持的数据类型 {type(data)}，无法打印表格。")
        if explain:
            logger.info(explain)
        else:
            None
        console.print(data_table(data_records))
        return None
    except FileNotFoundError:
        logger.error(f"文件 {data} 未找到，请检查文件路径。")
        return None
    except Exception as e:
        logger.error(f"读取文件时出现异常: {str(e)}")
        return None

def add_data_type(data):
    """
    此函数接收不同类型的数据（如 DataFrame、列表、字典），为其元素添加数据类型，并重新存储。
    对于 DataFrame，会将每列元素的数据类型添加到元素中；对于列表和字典，会为元素添加数据类型。

    参数：
    data：输入的数据，可以是 DataFrame、列表或字典。

    返回：
    存储了添加数据类型元素的新数据结构，与输入数据类型相同。
    """
    import pandas as pd
    if isinstance(data, pd.DataFrame):
        result_df = data.copy()
        for column in result_df.columns:
            result_df[column] = result_df[column].apply(lambda x: str(x) + ' (' + str(type(x).__name__) + ')')
        console_table(result_df)
    elif isinstance(data, list):
        result_list = []
        for item in data:
            if isinstance(item, dict):
                result_list.append(add_data_type(item))
            else:
                result_list.append(str(item) + ' (' + str(type(item).__name__) + ')')
        console_table(result_list)
    elif isinstance(data, dict):
        result_dict = {}
        for key, value in data.items():
            result_dict[key] = str(value) + ' (' + str(type(value).__name__) + ')'
        console_table(result_dict)
    else:
        raise ValueError("输入的数据类型不支持，仅支持 DataFrame、列表和字典。")
    
#-------------------------------------------------------------------------------------
# 自定义动画展示信息
#-------------------------------------------------------------------------------------

def rich_startup_info(input_explain=" ", info_upper="程序启动", info_below="欢迎使用" ):
    """
    自定义打印程序自动提示信息

    参数:
    input_explain: （可选）自定义内部提示信息，默认为" "。
    info_upper: （可选）自定义上边框信息。
    info_below: （可选）自定义下边框信息。
    """
    input_explain_text = Text(input_explain, style="bold yellow")
    centered_text = Align.center(input_explain_text)
    panel = Panel(centered_text, title=info_upper, subtitle=info_below, border_style="bold blue", expand=True)
    console.print(panel)



def art_startup_Info(text, font='random'):
    """
    参数：
    text: 需要打印的信息，由字符组成。
    font: 可选样式('block', 'slant', 'standard', 'banner', 'tarty', 'random')
    """
    result = text2art(text, font=font)
    console.print(result)